%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%% Explanation of what the file does
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 
% This function obtains a new guess for the path of utility hats in the RUV
% model with forward looking agents
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%% Inputs needed
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 
% 1. Matrix of sectoral price hats (PHma), size I*(S-1)*T
% 2. Final consumption shares (Amatrix), size I*(S-1)
% 3. Matrix of sectoral wage hats (Wmat), size I*(S-1)*T
% 4. Matrix of sectoral actual labor hats (Lmat), size I*(S-1)*T
% 5. Matrix of sectoral labor supplies (EllEra), size I*S*T
% 6. Mobility matrices (MuEra), size (I*S)*(I*S)*T
% 7. Discount factor (ba), scalar
% 8. Migration elasticity (ka), scalar
% 9. Sectoral elasticity (nu), scalar
% 10. Mobility matrices in baseline economy (MuBas), size (I*S)*(I*S)*T
% 11. The theta zero matrix for utility updating (TZ), size (I*S)*(I*S)
% 12. Labor supply matrix in baseline economy (EllBas), size I*(S+1)*T
% 13. Initial guess for U hat (Ug), size I*S*T
% 14. Degree of risk sharing (zrs), scalar
% 15. Matrix of sectoral actual labor dots in baseline (LDBas), size I*(S-1)*T
% 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%% Outputs produced
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 
% 1. New path of utility hats (UNew), size I*S*T
% 2. Delta risk sharing hats (Delrshat), size I*(S-1)*(T-1)

function [UNew,Delrshat] = NECO_Block3(PHma,Amatrix,Wmat,Lmat,EllEra,...
                           MuEra,ba,ka,nu,MuBas,TZ,EllBas,Ug,zrs,LDBas)

% Define some important variables

I                = size(EllEra,1);
S                = size(EllEra,2);
T                = size(EllEra,3);
EllRel           = EllEra(:,2:end,:);
EllBR            = EllBas(:,2:end,:);
EllDot           = EllRel(:,:,2:end)./EllRel(:,:,1:end-1);
EllDBas          = EllBR(:,:,2:end)./EllBR(:,:,1:end-1);
EllHat           = EllDot./EllDBas;
Pagg             = repmat(prod(PHma .^ repmat(Amatrix,1,1,T-1),2),1,S-1,1);

% Things for imperfect risk sharing

LEraB            = NaN(I,S-1,T);
LEraB(:,:,1)     = EllBR(:,:,1);
LEraB(:,:,2:end) = cumprod(LDBas,3).*repmat(EllBR(:,:,1),1,1,T-1);
unempprobB       = 1-LEraB./EllBR;
DeltarsB         = zrs.^(-unempprobB) .* ((1-unempprobB*zrs)./(1-unempprobB)).^(unempprobB-1);
DelrsdotB        = DeltarsB(:,:,2:end)./DeltarsB(:,:,1:end-1);

LDCfa            = Lmat.*LDBas;
LEraC            = NaN(I,S-1,T);
LEraC(:,:,1)     = EllRel(:,:,1);
LEraC(:,:,2:end) = cumprod(LDCfa,3).*repmat(EllRel(:,:,1),1,1,T-1);
unempprobC       = 1-LEraC./EllRel;
DeltarsC         = zrs.^(-unempprobC) .* ((1-unempprobC*zrs)./(1-unempprobC)).^(unempprobC-1);
DelrsdotC        = DeltarsC(:,:,2:end)./DeltarsC(:,:,1:end-1);

Delrshat         = DelrsdotC./DelrsdotB;

% Omegas and utilities

Omega            = ones(I,S,T-1);
Omega(:,2:end,:) = (Wmat.*Lmat)./(Pagg.*EllHat.*Delrshat);
UNew             = ones(I,S,T);

MuHashB          = NaN(I*S,I,T);
MuCondB          = NaN(size(MuBas));
for t=1:T
    MuHashB(:,:,t) = MuBas(:,:,t) * kron(eye(I),ones(S,1));
    MuCondB(:,:,t) = MuBas(:,:,t) ./ kron(MuHashB(:,:,t),ones(1,S));
end
MuCondB(isnan(MuCondB)) = 0;
MuHashDB         = MuHashB(:,:,2:end)./MuHashB(:,:,1:end-1);
MuHashDB(isnan(MuHashDB)) = 0;
MuCondDB         = MuCondB(:,:,2:end)./MuCondB(:,:,1:end-1);
MuCondDB(isnan(MuCondDB)) = 0;

% Run the backward induction loop where the future u dots are used to
% obtain the current u dots (the sectorial compensation omegas and mobility 
% matrices are also used)

for counter = 1:T-2
    t            = T-counter;
    UF           = Ug(:,:,t+1)';
    MuHash       = MuEra(:,:,t) * kron(eye(I),ones(S,1));
    MuCond       = MuEra(:,:,t) ./ kron(MuHash,ones(1,S));
    MuCond(isnan(MuCond)) = 0;
    Step1        = (repmat(UF(:)',I*S,1)).^(ba/nu) .* MuCond ...
                   .* MuCondDB(:,:,t);
    Step2        = (Step1 * kron(eye(I),ones(S,1))).^(nu/ka);
    Step3        = sum(MuHash .* MuHashDB(:,:,t) .* Step2,2).^ka;
    Step4        = reshape(Step3,S,I)';
    UNew(:,:,t)  = Omega(:,:,t) .* Step4;
end

% Period one is different

UF               = Ug(:,:,2)';
Inter            = TZ .* repmat(UF(:)',I*S,1).^(ba/nu);
Step2            = (Inter * kron(eye(I),ones(S,1))).^(nu/ka);
Step3            = Step2 .* MuHashB(:,:,2);
Step4            = sum(Step3,2) .^ ka;
Inner            = reshape(Step4,S,I)';
UNtemp           = Omega(:,:,1) .* Inner;
UNew(:,:,1)      = UNtemp./repmat(mean(UNtemp,2),1,S);

end
